#include <string>
#include <vector>
#include <windows.h>
#include <tchar.h>
#include <stdio.h>
#include <iostream>
#include <sstream>
#include <fstream>
#include <iomanip>
#include "opencv.hpp"
#include "CTruthDataPoint.h"
#include "DataPointManager.h"
#include "DataPoint.h"
#include "ImgTxtWriter.h"
#include "ColorPatchFeatExtractor.h"
#include "MorphologyFeatExtractor.h"
#include "SVMClassificationModule.h"
#include "svm.h"
/*******************************************************************/
using namespace cv;
using namespace std;
/*******************************************************************/
TTruthLabels ClassifyTestData( svm_model* model,TDataSet& data);
double GetError( TTruthLabels truth, TTruthLabels results ); 
svm_model* TrainSVM( TTrainingDataSet& dataset, svm_parameter& param); 
void WriteOutputLabels( CDataPointManager * mgr, TTruthLabels truth,
					   TTruthLabels results, TClassToNameMap nameMap,int &index, string name); 
/*******************************************************************/
int main(int argc, char* argv[])
{

	int sImgCount = 182;
	int nImgCount = 147;
	string sPath=".\\..\\..\\Data\\screws2\\";
	string baseName="datapoint";
	string ext=".jpg";
	// Load a bunch of sequential images from file
	CDataPointManager * screwTruthDataMgr = new CDataPointManager(baseName,sPath,ext,sImgCount/2,0,1);
	CDataPointManager * screwTestDataMgr = new CDataPointManager(baseName,sPath,ext,sImgCount/2,0,sImgCount/2);
	string nPath=".\\..\\..\\Data\\nuts\\";
	CDataPointManager * nailTruthDataMgr = new CDataPointManager(baseName,nPath,ext,nImgCount/2,0,1);
	CDataPointManager * nailTestDataMgr = new CDataPointManager(baseName,nPath,ext,nImgCount/2,0,nImgCount/2);
	
	CMorphologyFeatExtractor morphFeatureExtractor;

/*	svm_parameter ns_param;
	// default values
	ns_param.svm_type = C_SVC;
	ns_param.kernel_type = LINEAR;
	ns_param.degree = 3;
	ns_param.gamma = 0.2;
	ns_param.coef0 = 0;
	ns_param.nu = 0.5;
	ns_param.cache_size = 100;
	ns_param.C = 10;
	ns_param.eps = 1e-3;
	ns_param.p = 0.1;
	ns_param.shrinking = 1;
	ns_param.probability = 0;
	ns_param.nr_weight = 0;
	ns_param.weight_label = NULL;
	ns_param.weight = NULL;*/

	CSVMClassificationModule classifier;
	TClassToNameMap nameMap;
	nameMap[1] =  "bolt";
	nameMap[-1] = "nuts";
	string svmFile = "NailScrewClassifier";
	string svmPath = ".\\..\\..\\Data\\";
	TDataMgrVector trainData; 
	TTruthLabels trainLabels;
	trainData.push_back(screwTruthDataMgr);
	trainLabels.push_back(1.00);
	screwTestDataMgr->SetClassificationTruth(1.00);
	trainData.push_back(nailTruthDataMgr);
	trainLabels.push_back(-1.00); 
	nailTestDataMgr->SetClassificationTruth(-1.00);

	bool loadFromFile = true;
	if(  loadFromFile ) 
	{
		classifier.Load(svmFile,svmPath);
		classifier.SetFeatureExtractor(&morphFeatureExtractor);
	}
	else
	{
		classifier.SetClassLabels(nameMap);	
		classifier.SetFeatureExtractor(&morphFeatureExtractor);
		//classifier.SetParams(ns_param);


		classifier.BuildSVM(trainData,trainLabels); 
	}

	classifier.Classify(screwTestDataMgr);
	cout << "SCREW TEST ERROR: " << screwTestDataMgr->GetClassificationError() << endl;
	classifier.Classify(nailTestDataMgr);
	cout << "NAIL TEST ERROR: " << nailTestDataMgr->GetClassificationError() << endl;

	screwTruthDataMgr->SetClassificationTruth(1.00);
	classifier.Classify(screwTruthDataMgr);
	cout << "SCREW TRUTH ERROR: " << screwTruthDataMgr->GetClassificationError() << endl;
	nailTruthDataMgr->SetClassificationTruth(-1.00); 
	classifier.Classify(nailTruthDataMgr);
	cout << "NAIL TRUTH ERROR: " << nailTruthDataMgr->GetClassificationError() << endl;

	double errors = (screwTestDataMgr->GetClassificationError()+nailTestDataMgr->GetClassificationError())/2;
	cout << "TOTAL TEST ERROR: "  << errors << endl; 
	if( !loadFromFile )
	{
		if( classifier.Save(svmFile,svmPath) )
		{
			cout << "WIN!" << endl;
		}
		else
		{
			cout << "Epic Fail" << endl; 
		}
	}

/*	int imgCount = 120;
	string gPath=".\\..\\..\\Data\\grassland\\";
	
	baseName="image";
	ext=".jpg";
	
	// Load a bunch of sequential images from file
	CDataPointManager * grassTruthDataMgr = new CDataPointManager(baseName,gPath,ext,imgCount/2,3,1);
	CDataPointManager * grassTestDataMgr = new CDataPointManager(baseName,gPath,ext,imgCount/2,3,imgCount/2);
	string bPath=".\\..\\..\\Data\\beach\\";
	CDataPointManager * beachTruthDataMgr = new CDataPointManager(baseName,bPath,ext,imgCount/2,3,1);
	CDataPointManager * beachTestDataMgr = new CDataPointManager(baseName,bPath,ext,imgCount/2,3,imgCount/2);
	CColorPatchFeatExtractor colorFeatureExtractor;	svm_parameter param;

	// default values
	param.svm_type = C_SVC;
	param.kernel_type = LINEAR;
	param.degree = 3;
	param.gamma = 0;
	param.coef0 = 0;
	param.nu = 0.5;
	param.cache_size = 100;
	param.C = 1;
	param.eps = 1e-3;
	param.p = 0.1;
	param.shrinking = 1;
	param.probability = 0;
	param.nr_weight = 0;
	param.weight_label = NULL;
	param.weight = NULL;

	TTrainingDataSet grassBeachTrainSet;
	grassBeachTrainSet[-1] = beachTruthData;
	grassBeachTrainSet[1] = grassTruthData;
	svm_model* B_G_SVM = TrainSVM(grassBeachTrainSet, param);

	TTruthLabels grassTruthLabels;
	TTruthLabels beachTruthLabels;
	for( int k=0; k < imgCount/2; k++)
	{
		beachTruthLabels.push_back(-1); 
		grassTruthLabels.push_back(1); 
	}

	//TTruthLabels results;
	//double err = 0; 
	double bTestError = 0.00,bTruthError = 0.00,gTestError = 0.00,gTruthError = 0.00; 
	nameMap[-1] = "Beach Image";
	nameMap[1] = "Grass Image";
	outFile = "GrassBeachResult";
	index = 1;

	results = ClassifyTestData( B_G_SVM,beachTruthData);
	bTruthError = GetError(beachTruthLabels, results);
	cout << "Beach Truth Error Rate: " << bTruthError << endl;
	WriteOutputLabels( beachTruthDataMgr, beachTruthLabels, results, nameMap, index, outFile); 
	results.clear();

	results = ClassifyTestData( B_G_SVM,grassTruthData);
	gTruthError = GetError(grassTruthLabels, results);
	cout << "Grass Truth Error Rate: " << gTruthError << endl; 
	WriteOutputLabels( grassTruthDataMgr, grassTruthLabels, results, nameMap, index, outFile);
	results.clear();
	
	results = ClassifyTestData( B_G_SVM,beachTestData);
	bTestError = GetError(beachTruthLabels, results);
	cout << "Beach Test Error Rate: " << bTestError << endl; 
	WriteOutputLabels( beachTestDataMgr, beachTruthLabels, results, nameMap, index, outFile); 
	results.clear();

	results = ClassifyTestData( B_G_SVM,grassTestData);
	gTestError = GetError(grassTruthLabels, results);
	cout << "Grass Test Error Rate: " << gTestError << endl; 
	WriteOutputLabels( grassTestDataMgr, grassTruthLabels, results, nameMap, index, outFile);
	results.clear();

	// gah,.... this is dumb
	double bc = static_cast<double>(imgCount/2);
	double gc = static_cast<double>(imgCount/2);
	double bgTrainError = ((bTruthError*bc)+(gTruthError*gc))/(bc+gc);
	double bgTestError = ((bTestError*bc)+(gTestError*gc))/(bc+gc);
	cout << "TRAINING ERROR: " << bgTrainError << endl;
	cout << "TESTING ERROR: " << bgTestError << endl;
	svm_free_and_destroy_model(&B_G_SVM);
	svm_free_and_destroy_model(&N_S_SVM);
*/
	cout << "Bye!" << endl; 

	Sleep(100000);
	return 0;
	
}
/*******************************************************************/
TTruthLabels ClassifyTestData( svm_model* model,TDataSet& data)
{
	TTruthLabels retVal;
	int realFeats = static_cast<int>(data.begin()->size()); 
	int nfeats = realFeats+1; 
	svm_node* test_value = new svm_node[nfeats];
	TDataSetIter iter;
	for(  iter = data.begin(); iter != data.end(); ++iter)
	{
		for( int j=0;j<realFeats;++j)
		{
			test_value[j].value = iter->at(j);
			test_value[j].index = j+1;
		}
		test_value[nfeats].index = -1;

		double result = svm_predict(model, test_value);
		retVal.push_back(result);  
	}
	return retVal; 
}
/*******************************************************************/
double GetError( TTruthLabels truth, TTruthLabels results )
{
	int sz = static_cast<int>(truth.size());
	int count = 0; 
	for( int i=0; i < sz; i++ )
	{
		if( truth[i] != results[i] )
		{
			count++;
		}
	}
	return static_cast<double>(count)/static_cast<double>(sz); 
}
/*******************************************************************/
svm_model* TrainSVM( TTrainingDataSet& dataset, svm_parameter& param)
{
	svm_problem prob;
	TTrainingDataSetIter classIter;
	int samples = 0; 
	// get the total samples so we can allocate
	for( classIter = dataset.begin(); classIter != dataset.end(); ++classIter)
	{
		samples += static_cast<int>(((TDataSet)classIter->second).size());
	}
	TDataSet tempDataSet = dataset.begin()->second;
	int realFeats = static_cast<int>(tempDataSet.begin()->size());
	int nfeats = realFeats+1; // SVM's feature space appears to be plus one for some sort of terminator
	prob.l = samples; // size??
	prob.y = new double[samples]; // thinking this is label; 
	prob.x = new svm_node *[samples];
	svm_node *feature_space = new svm_node[nfeats * samples];
	TDataSetIter iter;
	int i=0;
	for(  classIter = dataset.begin(); classIter != dataset.end(); ++classIter )
	{
		tempDataSet = classIter->second;
		double label = classIter->first; 
		TDataSetIter iter; 
		for( iter = tempDataSet.begin(); iter != tempDataSet.end(); ++iter, ++i)
		{
			for( int j=0;j<realFeats;++j)
			{
				feature_space[nfeats*i+j].value = iter->at(j);
				feature_space[nfeats*i+j].index = j+1;
			}
			feature_space[nfeats*i+realFeats].index = -1;
			prob.x[i]=&feature_space[nfeats*i];
			prob.y[i] = static_cast<int>(label);
		}
	}

	svm_model* retVal = svm_train(&prob, &param);
	//delete [] prob.y; 
	//delete [] prob.x; 
	//delete [] feature_space;
	return retVal;
}
/*******************************************************************/
void WriteOutputLabels( CDataPointManager * mgr, TTruthLabels truth, TTruthLabels results, TClassToNameMap nameMap, int& index, string name)
{
	int sz = mgr->GetSize();
	CImgTxtWriter writer;
	for(int i=0; i<sz; i++)
	{
		mgr->GetDataPoint(i)->LoadSourceImage();
		if( results[i] == truth[i] )
		{
			writer.WriteText(nameMap[results[i]],mgr->GetDataPoint(i)->GetImage(),"GREEN",cvPoint(5,20));
		}
		else
		{
			writer.WriteText(nameMap[results[i]],mgr->GetDataPoint(i)->GetImage(),"RED",cvPoint(5,20));
			writer.WriteText(nameMap[truth[i]],mgr->GetDataPoint(i)->GetImage(),"GREEN",cvPoint(5,40));
		}
		ostringstream fname;
		fname << name << index << ".jpg";
		mgr->GetDataPoint(i)->SetOutFile(fname.str());
		mgr->GetDataPoint(i)->SaveImage(); 		
		index++;
	}
}
/*******************************************************************/